{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f0195c2-4a20-488e-8782-ca5a83488d0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.tools.waii import WaiiToolSpec\n",
    "\n",
    "waii_tool = WaiiToolSpec(\n",
    "    url=\"https://tweakit.waii.ai/api/\",\n",
    "    # API Key of Waii (not OpenAI API key)\n",
    "    api_key=\"3........\",\n",
    "    # Which database you want to use, you need add the db connection to Waii first\n",
    "    database_key=\"snowflake://....\",\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a79a9fa-e5ff-4242-99a2-08cc85e158a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "."
     ]
    },
    {
     "data": {
      "text/plain": [
       "'SELECT\\n    table_schema,\\n    table_name,\\n    COUNT(column_name) AS number_of_columns\\nFROM waii.information_schema.columns\\nGROUP BY\\n    table_schema,\\n    table_name\\nORDER BY\\n    table_schema,\\n    table_name\\n'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ".."
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>TABLE_SCHEMA</th>\n",
       "      <th>TABLE_NAME</th>\n",
       "      <th>NUMBER_OF_COLUMNS</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>BATTLE_DEATH</td>\n",
       "      <td>BATTLE</td>\n",
       "      <td>6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>BATTLE_DEATH</td>\n",
       "      <td>DEATH</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>BATTLE_DEATH</td>\n",
       "      <td>SHIP</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>CAR</td>\n",
       "      <td>CARS_DATA</td>\n",
       "      <td>8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>CAR</td>\n",
       "      <td>CAR_MAKERS</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>107</th>\n",
       "      <td>VOTER</td>\n",
       "      <td>VOTES</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>108</th>\n",
       "      <td>WORLD</td>\n",
       "      <td>CITY</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>109</th>\n",
       "      <td>WORLD</td>\n",
       "      <td>COUNTRY</td>\n",
       "      <td>15</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>110</th>\n",
       "      <td>WORLD</td>\n",
       "      <td>COUNTRYLANGUAGE</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>111</th>\n",
       "      <td>WORLD</td>\n",
       "      <td>SQLITE_SEQUENCE</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>112 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     TABLE_SCHEMA       TABLE_NAME  NUMBER_OF_COLUMNS\n",
       "0    BATTLE_DEATH           BATTLE                  6\n",
       "1    BATTLE_DEATH            DEATH                  5\n",
       "2    BATTLE_DEATH             SHIP                  7\n",
       "3             CAR        CARS_DATA                  8\n",
       "4             CAR       CAR_MAKERS                  4\n",
       "..            ...              ...                ...\n",
       "107         VOTER            VOTES                  5\n",
       "108         WORLD             CITY                  5\n",
       "109         WORLD          COUNTRY                 15\n",
       "110         WORLD  COUNTRYLANGUAGE                  4\n",
       "111         WORLD  SQLITE_SEQUENCE                  2\n",
       "\n",
       "[112 rows x 3 columns]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\"The table 'COLUMNS' contains the most columns. The top 5 tables with the number of columns are 'COLUMNS' with 43 columns, 'TABLES' with 25 columns, and the remaining tables have fewer than 25 columns.\""
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from llama_index import VectorStoreIndex\n",
    "\n",
    "# Use as Data Loader, load data to index and query it\n",
    "documents = waii_tool.load_data(\"Get all tables with their number of columns\")\n",
    "index = VectorStoreIndex.from_documents(documents).as_query_engine()\n",
    "\n",
    "index.query(\n",
    "    \"Which table contains most columns, tell me top 5 tables with number of columns?\"\n",
    ").response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b259d9cd-bbb8-4fff-a4ce-80fb0f3a1a10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use as tool, initialize it\n",
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "agent = FunctionAgent(\n",
    "    waii_tool.to_tool_list(), llm=OpenAI(model=\"gpt-4.1\"),\n",
    ")\n",
    "\n",
    "from llama_index.core.workflow import Context\n",
    "\n",
    "ctx = Context(agent)\n",
    "\n",
    "print(await agent.run(\"Give me top 3 countries with the most number of car factory\", ctx=ctx))\n",
    "print(await agent.run(\"What are the car factories of these countries\", ctx=ctx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90c2ba4d-6ac4-4cbb-93b0-e03a8d015042",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do performance analysis\n",
    "print(\n",
    "    await agent.run(\n",
    "        \"Give me top 3 longest running queries, include the complete query_id and their duration. And analyze performance of the first query\",\n",
    "        ctx=ctx,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47530eba-24be-42d9-b1a0-fa1af28934f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Diff two queries\n",
    "previous_query = \"\"\"\n",
    "SELECT\n",
    "    employee_id,\n",
    "    department,\n",
    "    salary,\n",
    "    AVG(salary) OVER (PARTITION BY department) AS department_avg_salary,\n",
    "    salary - AVG(salary) OVER (PARTITION BY department) AS diff_from_avg\n",
    "FROM\n",
    "    employees;\n",
    "\"\"\"\n",
    "current_query = \"\"\"\n",
    "SELECT\n",
    "    employee_id,\n",
    "    department,\n",
    "    salary,\n",
    "    MAX(salary) OVER (PARTITION BY department) AS department_max_salary,\n",
    "    salary - AVG(salary) OVER (PARTITION BY department) AS diff_from_avg\n",
    "FROM\n",
    "    employees;\n",
    "LIMIT 100;\n",
    "\"\"\"\n",
    "print(await agent.run(f\"tell me difference between {previous_query} and {current_query}\", ctx=ctx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a222df4-d00e-4fe8-be8e-9efdb43f1462",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Describe dataset\n",
    "print(await agent.run(\"Summarize the dataset\", ctx=ctx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6bdf837-241a-4637-a07e-a73fafd52a07",
   "metadata": {},
   "outputs": [],
   "source": [
    "q = \"\"\"\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql.functions import avg, lag, lead, round\n",
    "from pyspark.sql.window import Window\n",
    "\n",
    "spark = SparkSession.builder.appName(\"yearly_car_analysis\").getOrCreate()\n",
    "\n",
    "yearly_avg_hp = cars_data.groupBy(\"year\").agg(avg(\"horsepower\").alias(\"avg_horsepower\"))\n",
    "\n",
    "windowSpec = Window.orderBy(\"year\")\n",
    "\n",
    "yearly_comparisons = yearly_avg_hp.select(\n",
    "    \"year\",\n",
    "    \"avg_horsepower\",\n",
    "    lag(\"avg_horsepower\").over(windowSpec).alias(\"prev_year_hp\"),\n",
    "    lead(\"avg_horsepower\").over(windowSpec).alias(\"next_year_hp\")\n",
    ")\n",
    "\n",
    "final_result = yearly_comparisons.select(\n",
    "    \"year\",\n",
    "    \"avg_horsepower\",\n",
    "    round(\n",
    "        (yearly_comparisons.avg_horsepower - yearly_comparisons.prev_year_hp) / \n",
    "        yearly_comparisons.prev_year_hp * 100, 2\n",
    "    ).alias(\"percentage_diff_prev_year\"),\n",
    "    round(\n",
    "        (yearly_comparisons.next_year_hp - yearly_comparisons.avg_horsepower) / \n",
    "        yearly_comparisons.avg_horsepower * 100, 2\n",
    "    ).alias(\"percentage_diff_next_year\")\n",
    ").orderBy(\"year\")\n",
    "\n",
    "final_result.show()\n",
    "\"\"\"\n",
    "print(await agent.run(f\"translate this pyspark query {q}, to Snowflake\", ctx=ctx))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "myenv",
   "language": "python",
   "name": "myenv"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
